import torch
import itertools
from utils.image_pool import ImagePool
from .base_model import BaseModel
from .swagan import networks,conv2d_gradfix
from .swagan.non_leaking import AdaptiveAugment,augment
from .swagan.distributed import (
    get_rank,
    synchronize,
    reduce_loss_dict,
    reduce_sum,
    get_world_size,
)
from utils.util import SSIM,PSNR,normal,AverageMeter
import numpy as np
from  numpy.fft import fft2,fftshift
# from .lightcnn import network as lightcnn
import torch.nn.functional as F
from torch import autograd
from .cycle_gan_S012D_fft_model import NewAverageMeter
import math
import os

# import mypymath
# import time
from PIL import Image
from data.base_dataset import  get_transform
from torchvision import transforms
import random
def g_path_regularize(fake_img, latents, mean_path_length, decay=0.01):
    noise = torch.randn_like(fake_img) / math.sqrt(
        fake_img.shape[2] * fake_img.shape[3]
    )
    grad, = autograd.grad(
        outputs=(fake_img * noise).sum(), inputs=latents, create_graph=True
    )
    path_lengths = torch.sqrt(grad.pow(2).sum(2).mean(1))

    path_mean = mean_path_length + decay * (path_lengths.mean() - mean_path_length)

    path_penalty = (path_lengths - path_mean).pow(2).mean()

    return path_penalty, path_mean.detach(), path_lengths

def g_nonsaturating_loss(fake_pred):
    loss = F.softplus(-fake_pred).mean()

    return loss

def d_r1_loss(real_pred, real_img):
    with conv2d_gradfix.no_weight_gradients():
        grad_real, = autograd.grad(
            outputs=real_pred.sum(), inputs=real_img, create_graph=True
        )
    grad_penalty = grad_real.pow(2).reshape(grad_real.shape[0], -1).sum(1).mean()

    return grad_penalty


def d_logistic_loss(real_pred, fake_pred):
    real_loss = F.softplus(-real_pred)
    fake_loss = F.softplus(fake_pred)

    return real_loss.mean() + fake_loss.mean()

def make_noise(batch, latent_dim, n_noise, device):
    if n_noise == 1:
        return torch.randn(batch, latent_dim, device=device)

    noises = torch.randn(n_noise, batch, latent_dim, device=device).unbind(0)

    return noises

def mixing_noise(batch, latent_dim, prob, device):
    if prob > 0 and random.random() < prob:
        return make_noise(batch, latent_dim, 2, device)

    else:
        return [make_noise(batch, latent_dim, 1, device)]
    
def accumulate(model1, model2, decay=0.999):
    par1 = dict(model1.named_parameters())
    par2 = dict(model2.named_parameters())

    for k in par1.keys():
        par1[k].data.mul_(decay).add_(par2[k].data, alpha=1 - decay)

class SWAGANModel(BaseModel):
    """
    This class implements the CycleGAN model, for learning image-to-image translation without paired data.

    The model training requires '--dataset_mode unaligned' dataset.
    By default, it uses a '--netG resnet_9blocks' ResNet generator,
    a '--netD basic' discriminator (PatchGAN introduced by pix2pix),
    and a least-square GANs objective ('--gan_mode lsgan').

    CycleGAN paper: https://arxiv.org/pdf/1703.10593.pdf
    """
    @staticmethod
    def modify_commandline_options(parser, is_train=True):
        """Add new dataset-specific options, and rewrite default values for existing options.

        Parameters:
            parser          -- original option parser
            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.

        Returns:
            the modified parser.

        For CycleGAN, in addition to GAN losses, we introduce lambda_A, lambda_B, and lambda_identity for the following losses.
        A (source domain), B (target domain).
        Generators: G_A: A -> B; G_B: B -> A.
        Discriminators: D_A: G_A(A) vs. B; D_B: G_B(B) vs. A.
        Forward cycle loss:  lambda_A * ||G_B(G_A(A)) - A|| (Eqn. (2) in the paper)
        Backward cycle loss: lambda_B * ||G_A(G_B(B)) - B|| (Eqn. (2) in the paper)
        Identity loss (optional): lambda_identity * (||G_A(B) - B|| * lambda_B + ||G_B(A) - A|| * lambda_A) (Sec 5.2 "Photo generation from paintings" in the paper)
        Dropout is not used in the original CycleGAN paper.
        """
        
        # parser.set_defaults(frequency=True)
        # parser.add_argument('--num_classes', type=int, default=3, help='the architecture situation of model')
        # parser.add_argument('--pretrained_model', type=str, default="lightcnn", help='the architecture of pretrained model')
        # parser.add_argument('--pretrained_path', type=str, default="lightcnn", help='the path of pretrained model')
        # parser.add_argument('--pretrained_ndf', type=int, default=128, help='the number of features of output ')
        # parser.add_argument('--pretrained_output_nc', type=int, default=256, help='the number of features of output ')

        parser.set_defaults(no_dropout=True)  # default CycleGAN did not use dropout
        parser.add_argument('--inputs', type=str, default='0', help='the moudal of input data ')
        parser.add_argument('--outputs', type=str, default='1', help='the moudal ofoutput data ')
        # parser.add_argument('--bins', type=int, default=16, help='the number of features of output ')
        # parser.add_argument('--cal_mode', type=str, default='linear', help='the mode of caculating')
        # parser.add_argument('--slope', type=int, default=0, help='slope parameter')
        # parser.add_argument('--frequency', type=bool, default=True, help='the mode of caculating')
        # if is_train:
        parser.add_argument('--sizes', type=int, default=256, help='image sizes for the model')
        parser.add_argument('--latent', type=int, default=512, help='latent')
        parser.add_argument('--n_mlp', type=int, default=8, help='n_mlp')
        
        parser.add_argument("--r1", type=float, default=10, help="weight of the r1 regularization")
        parser.add_argument("--path_regularize",type=float,default=2,help="weight of the path length regularization")
        parser.add_argument("--path_batch_shrink",type=int,default=2,help="batch size reducing factor for the path length regularization (reduce memory consumption)",)
        parser.add_argument("--d_reg_every",type=int,default=16,help="interval of the applying r1 regularization")
        parser.add_argument("--g_reg_every",type=int,default=4,help="interval of the applying path length regularization",)
        parser.add_argument("--mixing", type=float, default=0.9, help="probability of latent code mixing")
        parser.add_argument("--channel_multiplier",type=int,default=2,help="channel multiplier factor for the model. config-f = 2, else = 1",)
        # parser.add_argument("--wandb", action="store_true", help="use weights and biases logging")
        parser.add_argument("--local_rank", type=int, default=0, help="local rank for distributed training")
        parser.add_argument("--augment", action="store_true", help="apply non leaking augmentation")
        parser.add_argument("--augment_p",type=float,default=0,help="probability of applying augmentation. 0 = use adaptive augmentation",)
        parser.add_argument("--ada_target",type=float,default=0.6,help="target augmentation probability for adaptive augmentation",)
        parser.add_argument("--ada_length",type=int,default=500 * 1000,help="target duraing to reach augmentation probability for adaptive augmentation",)
        parser.add_argument("--ada_every",type=int,default=256,help="probability update interval of the adaptive augmentation",)
        
        
        parser.add_argument('--lambda_A', type=float, default=10.0, help='weight for cycle loss (A -> B -> A)')
        parser.add_argument('--lambda_B', type=float, default=10.0, help='weight for cycle loss (B -> A -> B)')
        parser.add_argument('--lambda_fre', type=float, default=0, help='weight for preTrained Discriminator loss')
        parser.add_argument('--lambda_identity', type=float, default=0.5, help='use identity mapping. Setting lambda_identity other than 0 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set lambda_identity = 0.1')
        return parser
    
    def display_opt_init(opt):
        # parser.add_argument('--inputs', type=str, default='img', help='the moudal of input data ')
        opt.__setattr__('inputs','S0')
        opt.__setattr__('outputs','SD')
        # opt.__setattr__('cal_mode','linear')
        opt.__setattr__('lambda_fre',0)
        return opt
    def display_init(self,_dict):
        
        transform = get_transform(self.opt, grayscale=(self.opt.input_nc == 1))
        # pil_img = Image.fromarray(np.uint8(_dict['img']))
        # _dict[self.opt.inputs] = transform(pil_img)
        # c,h,w = _dict[self.opt.inputs].shape
        # _dict[self.opt.inputs] = _dict[self.opt.inputs].reshape(1,c,h,w)
        _dict[self.opt.inputs+'_paths'] = "ZedCamera"
        for b in _dict['boxes']:
            face = transform(Image.fromarray(np.uint8(_dict['img'][b[1]:b[3],b[0]:b[2],:])))
            c,h,w = face.shape
            if self.opt.inputs in _dict.keys():
                _dict[self.opt.inputs] = torch.cat((_dict[self.opt.inputs],face.unsqueeze(0)),0)
            else:
                _dict[self.opt.inputs] = face.unsqueeze(0)
        return _dict

    def __init__(self, opt):
        """Initialize the CycleGAN class.

        Parameters:
            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
        """
        
        BaseModel.__init__(self, opt)
        # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
        # self.loss_names = ['D_A', 'G_A', 'cycle_A', 'idt_A', 'D_B', 'G_B', 'cycle_B', 'idt_B']
        if self.opt.controller == 'test':
            self.opt.lambda_fre = 0
        self.opt.frequency = not self.opt.lambda_fre == 0
        self.A = self.opt.inputs
        self.B = self.opt.outputs
        self.cur_epoch = 0
        # self.last_val_best = 0
        self.visuale_param()

        # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>.
        if self.isTrain:
            # self.model_names = ['G_A', 'G_B', 'D_A', 'D_B']
            self.model_names = ['G', 'G_ema','D', ]
        else:  # during test time, only load Gs
            self.model_names = ['G_SAB',]

        # define networks (both Generators and discriminators)
        # The naming is different from those used in the paper.
        # Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
        self.netG = networks.define_G(opt.sizes , opt.latent , opt.n_mlp, opt.channel_multiplier, opt.init_type, opt.init_gain, self.gpu_ids)
        self.netG_ema = networks.define_G(opt.sizes , opt.latent , opt.n_mlp, opt.channel_multiplier, opt.init_type, opt.init_gain, self.gpu_ids)
        self.netG_ema.eval()
        accumulate(self.netG_ema, self.netG, 0)
        
        # self.netG_SBA = networks.define_G(opt.size , opt.latent , opt.n_mlp, opt.channel_multiplier, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
        # self.netG_SBA_ema = networks.define_G(opt.size , opt.latent , opt.n_mlp, opt.channel_multiplier, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
        # self.netG_SBA_ema.eval()
        # accumulate(self.netG_SBA_ema, self.netG_SBA, 0)

        if self.isTrain:  # define discriminators
            self.netD = networks.define_D(opt.sizes, opt.channel_multiplier, opt.init_type, opt.init_gain, self.gpu_ids)
            
            # self.netD_SB = networks.define_D(opt.sizes, opt.channel_multiplier, opt.init_type, opt.init_gain, self.gpu_ids)
            # self.netD_SA = networks.define_D(opt.sizes, opt.channel_multiplier, opt.init_type, opt.init_gain, self.gpu_ids)
            # self.netD_light.define_lightcnn(opt)

            # self.netD_SA2 = networks.define_D(opt.input_nc, opt.ndf, opt.netD,    # opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
            # if self.opt.val:
            #     self.last_val_best = 0
        # self.criterionPredD = torch.nn.L1Loss()
        
      
        
        if self.isTrain:
            if self.opt.continue_train:
                self.epoch = self.opt.epoch_count
            # self.nepoch = self.opt.n_epochs + self.opt.n_epochs_decay
            # self.fac = self.epoch/(self.nepoch+1)
            if opt.lambda_identity > 0.0:  # only works when input and output images have the same number of channels
                assert(opt.input_nc == opt.output_nc)
            # self.fake_SA_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            # self.fake_SB_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            # self.fake_S2_SA_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
            
            # define loss functions
            # self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)  # define GAN loss.
            # self.criterionCycle = torch.nn.L1Loss()
            # self.criterionIdt = torch.nn.L1Loss()
        
        
        # if self.isTrain:
            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
            g_reg_ratio = opt.g_reg_every / (opt.g_reg_every + 1)
            d_reg_ratio = opt.d_reg_every / (opt.d_reg_every + 1)
            self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr*g_reg_ratio, betas=(opt.beta1, 0.999**g_reg_ratio))
            self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr*d_reg_ratio, betas=(opt.beta1, 0.999**g_reg_ratio))
            self.optimizers.append(self.optimizer_G)
            self.optimizers.append(self.optimizer_D)
            
            self.loss_D = 0
            self.loss_r1 = torch.tensor(0.0, device=self.device)
            self.loss_G = 0
            self.loss_path = torch.tensor(0.0, device=self.device)
            self.loss_path_lengths = torch.tensor(0.0, device=self.device)
            self.loss_mean_path_length_avg = 0
            self.mean_path_length = 0
            self.accum = 0.5 ** (32 / (10 * 1000))
            self.ada_aug_p = self.opt.augment_p if self.opt.augment_p > 0 else 0.0
            self.loss_r_t_stat = 0
            
            # _model = opt.model
            # _isTrain = opt.isTrain
            # _ndf = opt.ndf
            # _outpu_nc = opt.output_nc
            # opt.model = opt.pretrained_model
            # opt.isTrain = False
            # opt.ndf = opt.pretrained_ndf
            # opt.output_nc = opt.pretrained_output_nc
            # self.netD_light = create_model(opt)      # create a model given opt.model and other options
            # self.netD_light.setup(opt)
            # self.netD_light.load_pretrained_networks(opt.pretrained_path)
            # opt.model = _model
            # opt.isTrain = _isTrain
            # opt.ndf = _ndf
            # opt.output_nc = _outpu_nc
    def visuale_param(self):
        if self.isTrain:
            self.loss_names = ["G","D","r1","path","real_score","fake_score", "path_length","mean_path_length_avg","ada_aug_p"]
            self.count_G=AverageMeter()
        else:
            self.loss_names = ["SB_SSIM","SAB_SSIM","SB_PSNR","SAB_PSNR"]
            self.loss_names += ["mean_SB_PSNR","mean_SAB_PSNR"]
            self.count_mean_SB_PSNR=AverageMeter()
            self.count_mean_SAB_PSNR=AverageMeter()
        
            self.loss_names += ["mean_SB_SSIM","mean_SAB_SSIM"]
            self.count_mean_SB_SSIM=NewAverageMeter()
            self.count_mean_SAB_SSIM=AverageMeter()
            # if self.opt.controller == 'test':
            # self.count_mean_S012D_SSIM=NewAverageMeter()
        # self.count_ppaGSBA=AverageMeter()
        # self.count_ppaGSAB=AverageMeter()
        # self.count_ppaRSBA=AverageMeter()
        # self.count_ppaRSAB=AverageMeter()
        if self.opt.lambda_fre != 0:
            self.loss_names +=["mean_Frequency"]
            self.count_Frequency=AverageMeter()
        
        visual_names_S1A = ['real_SA', 'fake_SBA', ]
        visual_names_S1B = ['real_SB', 'fake_SAB']
        # if self.isTrain:
        visual_names_S1B += ['rec_SAB', 'rec_SBA']
        # visual_names_S1A += ['real_SA_h', 'fake_SBA_h', 'rec_SBA_h']
        # visual_names_S1B += ['real_SB_h', 'fake_SAB_h', 'rec_SAB_h']
        
        self.visual_names = visual_names_S1A + visual_names_S1B  # combine visualizations for A and B
        pass

    def set_input(self, input):
        """Unpack input data from the dataloader and perform necessary pre-processing steps.

        Parameters:
            input (dict): include the data itself and its metadata information.

        The option 'direction' can be used to swap domain A and domain B.
        """
        self.cur_epoch += 1
        # AtoB = self.opt.direction == 'AtoB'
        # self.real_A = input['A' if AtoB else 'B'].to(self.device)
        # self.real_B = input['B' if AtoB else 'A']self.device)
        # self.image_paths = input['A_paths' if AtoB else 'B_paths']
        if self.A in input.keys():
            self.real_img = input[self.A].to(self.device)
            self.image_paths = input[self.A+'_paths']
        else:
            self.real_img = None
            self.image_paths = None
            
        if self.B in input.keys():
            self.fake_img = input[self.B].to(self.device)
        else:
            self.fake_img = None
        # self.real_S2 = input['S2'].to(self.device)
        
        if "isTrain" in input.keys():
            if self.isTrain != input["isTrain"]:
                self.isTrain = input["isTrain"]
                self.visuale_param()

    def forward_D(self):
        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
        # self.fake_B = self.netG_A(self.real_A)  # G_A(A)
        # self.rec_A = self.netG_B(self.fake_B)   # G_B(G_A(A))
        # self.fake_A = self.netG_B(self.real_B)  # G_B(B)
        # self.rec_B = self.netG_A(self.fake_A)   # G_A(G_B(B))
        
        noise = mixing_noise(self.opt.batch_size, self.opt.latent, self.opt.mixing, self.device)
        self.fake_img, _ = self.netG(noise)
        
        if self.opt.augment:
            self.real_img_aug, _ = augment(self.real_img, self.ada_aug_p)
            self.fake_img, _ = augment(self.fake_img, self.ada_aug_p)

        else:
            self.real_img_aug = self.real_img
        
        self.fake_pred = self.netD(self.fake_img)
        self.real_pred = self.netD(self.real_img_aug)
        
        # if self.real_SA is not None:
        #     self.fake_SAB = self.netG_SAB(self.real_SA)
        #     self.rec_SBA = self.netG_SBA(self.fake_SAB)
        #     # if self.isTrain:
        #     if self.dualGen:
        #         self.fake_SBA = self.netG_SBA(self.real_SB)
        #         self.rec_SAB = self.netG_SAB(self.fake_SBA)
    def forward(self):
        noise = mixing_noise(self.opt.batch_size, self.opt.latent, self.opt.mixing, self.device)
        self.fake_img, _ = self.netG(noise)

        if self.opt.augment:
            self.fake_img, _ = augment(self.fake_img, self.ada_aug_p)

        self.fake_pred = self.netD(self.fake_img)
        
    def backward_D(self,):
        """Calculate GAN loss for the discriminator

        Parameters:
            netD (network)      -- the discriminator D
            real (tensor array) -- real images
            fake (tensor array) -- images generated by a generator

        Return the discriminator loss.
        We also call loss_D.backward() to calculate the gradients.
        """
        d_loss = d_logistic_loss(self.real_pred, self.fake_pred)

        
        # self.optimizer_D.zero_grad() 
        self.netD.zero_grad()
        d_loss.backward()
        self.optimizer_D.step()
        self.loss_D=d_loss.mean().item()
        self.loss_real_score = self.real_pred.mean().mean().item()
        self.loss_fake_score = self.fake_pred.mean().mean().item()
        if self.opt.augment and self.opt.augment_p == 0:
            self.ada_aug_p = self.ada_augment.tune(self.real_pred)
            self.loss_r_t_stat = self.ada_augment.r_t_stat

        d_regularize = self.cur_epoch % self.opt.d_reg_every == 0

        if d_regularize:
            self.real_img.requires_grad = True

            if self.opt.augment:
                self.real_img_aug, _ = augment(self.real_img, self.ada_aug_p)

            else:
                self.real_img_aug = self.real_img

            self.real_pred = self.netD(self.real_img_aug)
            self.loss_r1 = d_r1_loss(self.real_pred, self.real_img).mean().item()

            self.netD.zero_grad()
            (self.opt.r1 / 2 * self.loss_r1 * self.opt.d_reg_every + 0 * self.real_pred[0]).backward()
            self.optimizer_D.step()
        self.loss_ada_aug_p = self.ada_aug_p
    
    def backward_G(self):
        """Calculate the loss for generators G_A and G_B"""
        g_loss = g_nonsaturating_loss(self.fake_pred)

        self.netG.zero_grad()
        g_loss.backward()
        self.optimizer_G.step()
        self.loss_G = g_loss.mean().item()
        
        g_regularize = self.cur_epoch % self.opt.g_reg_every == 0

        if g_regularize:
            path_batch_size = max(1, self.opt.batch_size // self.opt.path_batch_shrink)
            noise = mixing_noise(path_batch_size, self.opt.latent, self.opt.mixing, self.device)
            fake_img, latents = self.netG(noise, return_latents=True)

            self.loss_path, self.mean_path_length, self.loss_path_lengths = g_path_regularize(
                fake_img, latents, self.mean_path_length
            )

            self.netG.zero_grad()
            weighted_path_loss = self.opt.path_regularize * self.opt.g_reg_every * self.loss_path

            if self.opt.path_batch_shrink:
                weighted_path_loss += 0 * fake_img[0, 0, 0, 0]

            weighted_path_loss.backward()

            self.optimizer_G.step()

            self.loss_mean_path_length_avg = (
                reduce_sum(self.mean_path_length).item() / get_world_size()
            )
            self.loss_path = self.loss_path.mean().item()
            self.loss_path_length = self.loss_path_lengths.mean().mean().item()

    def optimize_parameters(self):
        """Calculate losses, gradients, and update network weights; called in every training iteration"""
        # forward
        
        
        
        

        
        if self.opt.augment and self.opt.augment_p == 0:
            self.ada_augment = AdaptiveAugment(self.opt.ada_target, self.opt.ada_length, 8, self.device)

        # sample_z = torch.randn(self.opt.n_sample, self.opt.latent, device=self.device)
        
        
        # requires_grad(generator, False)
        self.set_requires_grad([self.netG], False)
        self.set_requires_grad([self.netD], True)
        self.forward_D()
        # loss_dict["d"] = d_loss
        
        self.backward_D()
       

        # loss_dict["r1"] = r1_loss
        

        # requires_grad(generator, True)
        # requires_grad(discriminator, False)
        self.set_requires_grad([self.netG], True)
        self.set_requires_grad([self.netD], False)

        self.forward()
        self.backward_G()

        # loss_dict["path"] = path_loss
        
        
        # loss_dict["path_length"] = path_lengths.mean()

        accumulate(self.netG_ema, self.netG, self.accum)
        
    # def cal_score(self):
    #     if not self.isTrain:
    #         self.backward_G() 
        # self.fake_SAB = recoverFslope(self.fake_SAB,1,-1)
        # self.real_SA = recoverFslope(self.real_SA,1,-1)
       
        
        # if not self.isTrain:
           
        # self.fake_S02 = normal(self.fake_S02,1,0)  
        # gt_SA = image2gt(self.real_SA, self.opt.bins,self.opt.cal_mode)
        # pred_SBA = image2gt(self.fake_SBA, self.opt.bins,self.opt.cal_mode)

        # rec_SBA = image2gt(self.rec_SBA, self.opt.bins,self.opt.cal_mode)
        
        # gt_SB = image2gt(self.real_SB, self.opt.bins,self.opt.cal_mode)
        # pred_SB = image2gt(self.fake_SAB, self.opt.bins,self.opt.cal_mode)
        # rec_SB = image2gt(self.rec_SAB, self.opt.bins,self.opt.cal_mode)
        
        # self.real_SA_h = gt_SA[0,:,:]/np.max(gt_SA)*255
        # self.fake_SBA_h = pred_SBA[0,:,:]/np.max(pred_SBA)*255
        # self.rec_SBA_h = rec_SBA[0,:,:]/np.max(rec_SBA)*255
        # self.real_SB_h = gt_SB[0,:,:]/np.max(gt_SB)*255
        # self.fake_SAB_h = pred_SB[0,:,:]/np.max(pred_SB)*255
        # self.rec_SAB_h = rec_SB[0,:,:]/np.max(rec_SB)*255

        # self.count_ppaGSBA.update_np(per_pix_acc(gt_SA,pred_SBA))
        # self.count_ppaGSAB.update_np(per_pix_acc(gt_SB,pred_SB))
        # self.count_ppaRSBA.update_np(per_pix_acc(gt_SA,rec_SBA))
        # self.count_ppaRSBA.update_np(per_pix_acc(gt_SB,rec_SB))
       
        # self.loss_per_pix_acc_G_SBA = self.count_ppaGSBA.avg  # per_pix_acc(gt_SA,pred_SBA)
        # self.loss_per_pix_acc_G_SAB = self.count_ppaGSAB.avg  # per_pix_acc(gt_SB,pred_SB)
        
        # self.loss_per_pix_acc_rec_SBA = self.count_ppaRSBA.avg  # per_pix_acc(gt_SA,rec_SBA)
        # self.loss_per_pix_acc_rec_SAB = self.count_ppaRSBA.avg  # per_pix_acc(gt_SB,rec_SB)

        pass
    def iter_end(self):
        
        if self.opt.controller == 'test' or not self.isTrain:
            if self.dualGen:
                # self.real_SB = recoverFslope(self.real_SB,1,-1)

                self.loss_SB_PSNR = PSNR(self.fake_img,self.real_img)
                self.loss_SAB_PSNR = PSNR(self.real_SA,self.real_SB)
                self.loss_SB_SSIM = SSIM(self.fake_SAB,self.real_SB)
                self.loss_SAB_SSIM = SSIM(self.real_SA,self.real_SB)
                
                self.count_mean_SB_PSNR.update(self.loss_SB_PSNR)
                self.count_mean_SAB_PSNR.update(self.loss_SAB_PSNR)
                self.count_mean_SB_SSIM.update(self.loss_SB_SSIM, self.opt.batch_size, paths = self.image_paths)
                self.count_mean_SAB_SSIM.update(self.loss_SAB_SSIM)
                # self.count_mean_S012D_SSIM.update(self.loss_SB_SSIM, self.opt.batch_size, paths = self.image_paths)
                
                self.loss_mean_SB_PSNR = self.count_mean_SB_PSNR.avg
                self.loss_mean_SAB_PSNR = self.count_mean_SAB_PSNR.avg
                self.loss_mean_SB_SSIM = self.count_mean_SB_SSIM.avg
                self.loss_mean_SAB_SSIM = self.count_mean_SAB_SSIM.avg
                
                # self.loss_S012D_SSIM = self.loss_S01_SSIM + self.loss_S02_SSIM + self.loss_S0D_SSIM
                
                self.TOP100p_mean_SB_SSIM = self.count_mean_SB_SSIM.avg
        
    def epoch_end(self):
        if self.opt.lambda_fre != 0:
            self.count_Frequency.reset()
        if self.isTrain:
            # self.count_ppaGSBA.reset()
            # self.count_ppaGSAB.reset()
            # self.count_ppaRSBA.reset()
            # self.count_ppaRSBA.reset()
            self.count_G.reset()
        else:
            self.count_mean_SB_SSIM.statistic()
            if self.opt.controller == 'test':
                self.count_mean_SB_SSIM.saveTop10p(os.path.join(os.path.abspath('.'),'results',self.opt.name+'_v'+self.opt.version,self.opt.phase+'_'+self.opt.epoch),'Top100pImgPaths.txt')
            self.TOP10p_mean_SB_SSIM = self.count_mean_SB_SSIM.Top10p
            self.TOP50p_mean_SB_SSIM = self.count_mean_SB_SSIM.Top50p


            self.count_mean_SB_PSNR.reset()
            self.count_mean_SAB_PSNR.reset()
        
            self.count_mean_SB_SSIM.reset()
            self.count_mean_SAB_SSIM.reset()
        
        # if self.opt.isTrain:
        #     self.epoch+=1
            # self.fac = self.epoch/(self.nepoch+1)
    def val_save(self,param):
        self.val_best = getattr(self,param)
        if self.last_val_best[param] < self.val_best:
            self.last_val_best[param] = self.val_best
            return True
        return False

    def display_proc(self,_dict):
        """
            为了展示控制器准备的函数接口，用户可以自己定义
        """
        # image)
        tran = transforms.ToPILImage()
        if len(_dict['boxes'])>0:
            B,C,H,W = self.fake_SAB.shape
            for i,b in enumerate(_dict['boxes']):
                face = tran(self.fake_SAB[i,:,:,:].contiguous().cpu().detach().squeeze(0))
                h = b[3] - b[1]
                w = b[2] - b[0]
                
                _dict['box_img'][b[1]:b[3],b[0]:b[2],0] = np.array(face.resize((w,h)))
                _dict['box_img'][b[1]:b[3],b[0]:b[2],1] = np.array(face.resize((w,h)))
                _dict['box_img'][b[1]:b[3],b[0]:b[2],2] = np.array(face.resize((w,h)))
                _dict['SD'] = self.fake_SAB
                
            pass
            
        return _dict

class AverageMeter(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count
    
    def update_np(self,val):
        self.sum = np.sum(val,0)
        self.count += val.shape[0]
        self.avg = self.sum / self.count